import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['font.size'] = 12

def generate_housing_data():
    features = ['Median Income', 'House Age', 'Rooms', 'Bedrooms', 'Population', 'Households', 'Latitude', 'Longitude']
    original = [4.0, 28.0, 1450, 290, 1350, 480, 34.2, -118.3]
    fedavg_recon = [4.2, 28.5, 1500, 300, 1400, 500, 34.5, -118.2]
    nrfl_recon = [3.8, 25.0, 1200, 280, 1100, 450, 34.0, -118.0]
    return features, original, fedavg_recon, nrfl_recon

def plot_feature_reconstruction():
    features, original, fedavg, nrfl = generate_housing_data()
    
    # 创建画布和GridSpec布局（2行1列，高度比3:1）
    plt.figure(figsize=(5, 5))
    gs = GridSpec(2, 1, height_ratios=[3, 1])  # 关键修复点[4](@ref)[9](@ref)
    
    # 主图：特征重建对比（第1行）
    ax1 = plt.subplot(gs[0, 0])  # 修正索引语法[2](@ref)[7](@ref)
    x = np.arange(len(features))
    width = 0.25
    
    # 绘制柱状图
    rects1 = ax1.bar(x - width, original, width, label='Original Data', color='#4daf4a', alpha=0.8)
    rects2 = ax1.bar(x, fedavg, width, label='FedAvg Reconstruction', color='#e41a1c', alpha=0.8)
    rects3 = ax1.bar(x + width, nrfl, width, label='NRFL Reconstruction', color='#377eb8', alpha=0.8)
    
    # 添加标注
    ax1.set_ylabel('Feature Value', fontsize=16)
#     ax1.set_title('(a) Feature Reconstruction Comparison on California Housing', fontsize=16)
    ax1.set_xticks(x)
    ax1.set_xticklabels(features, rotation=20, ha='right')
    ax1.legend(loc='upper right', fontsize=12)
#     ax1.grid(axis='y', alpha=0.3)
    
    # 添加PSNR标注
#     ax1.text(0.5, 1600, 'FedAvg PSNR: 32.6 dB (High Risk)', fontsize=12, color='#e41a1c', ha='center')
#     ax1.text(6.5, 1600, 'NRFL PSNR: 14.8 dB (Low Risk)', fontsize=12, color='#377eb8', ha='center')
    
    # 子图：标签恢复率对比（第2行）
    ax2 = plt.subplot(gs[1, 0])  # 修正索引语法[2](@ref)[7](@ref)
    labels = ['FedAvg', 'NRFL']
    recovery_rates = [85.0, 41.3]
    colors = ['#e41a1c', '#377eb8']
    
    bars = ax2.bar(labels, recovery_rates, color=colors, alpha=0.8)
    ax2.set_ylabel('Label Recovery Rate', fontsize=16)
    ax2.set_ylim(0, 100)
#     ax2.grid(axis='y', alpha=0.3)
    
    # 添加数据标签
    for bar in bars:
        height = bar.get_height()
        ax2.text(bar.get_x() + bar.get_width()/2, height+2, f'{height}%', ha='center', va='bottom', fontsize=12)
    
    # 调整布局
    plt.tight_layout()
    plt.subplots_adjust(top=1.3, hspace=0.3)  # 增加行间距[9](@ref)
    plt.savefig('Fig7.png', dpi=800, bbox_inches='tight')
    plt.show()

# 执行绘图
plot_feature_reconstruction()
